open-standards/wai/wai-rs/src/world_model.rs

//! `wai.world.model` — **generative-world transport** with replay-equivalence on a
//! signed action trace + per-frame provenance.
//!
//! The interactive-world-model field (Genie / Oasis / Matrix / Hunyuan-World class)
//! transports *pixels*: a GPU runs the world model and streams generated video to a
//! thin client, the client sends controls back. Nothing about that session is
//! reproducible, auditable, or energy-accounted — the generated world is ephemeral.
//!
//! WAI flips the wire, the same way [`crate::world`] does for deterministic worlds:
//! the canonical artifact is the **action trace**, not the pixels. Each control is a
//! signed [`CausalOp`] (the world spec already says "an input event is an op emitted
//! by one actor"); the trace replays — under canonical total order — to one
//! **session hash** on every machine (*replay-equivalence*). A declared `model_id`
//! pins WHICH world model + seed turns that trace into frames; the model itself is
//! sink-supplied presentation. The generated frames are vouched for by a signed,
//! energy-accounted provenance receipt ([`crate::provenance::AssetReceipt`]) —
//! content-addressed frames, joules per frame.
//!
//! So a `wai.world.model` session is portable, signed, metered, and replayable even
//! when the renderer is a float model. When the renderer is an INTEGER-exact world
//! model (the `int_*` family), the frames themselves hash-converge too — a
//! generative interactive world you can put in a receipt.

use ed25519_dalek::SigningKey;

use crate::provenance::{AssetReceipt, ProvStep};
use crate::world::{CausalOp, OpId};

/// Domain tag for the session hash-chain (distinct from `wai:world-state`: this
/// chains the ACTION TRACE, not a simulated state).
pub const DOMAIN_MODEL_SESSION: &[u8] = b"wai:world-model-session\x01";

/// `CausalOp::mutation_kind` for an action input (move/look/buttons + free payload).
pub const MUT_ACTION: u16 = 1;

/// The declared world model + generation parameters. The renderer is sink-supplied;
/// this header pins WHAT (model + seed + frame geometry) turns the action trace into
/// frames, so it is bound into the session hash.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ModelHeader {
    /// Content hash / DID key-bytes of the world model (the renderer identity).
    pub model_id: [u8; 32],
    /// Generation seed — pins the model's stochastic path for replay.
    pub seed: u64,
    pub width: u32,
    pub height: u32,
    pub fps: u32,
}

impl ModelHeader {
    /// Canonical bytes bound into the session hash.
    pub fn preimage(&self) -> Vec<u8> {
        let mut o = Vec::with_capacity(52);
        o.extend_from_slice(&self.model_id);
        o.extend_from_slice(&self.seed.to_le_bytes());
        o.extend_from_slice(&self.width.to_le_bytes());
        o.extend_from_slice(&self.height.to_le_bytes());
        o.extend_from_slice(&self.fps.to_le_bytes());
        o
    }
}

/// One decoded control fed to the world model at a tick.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Action {
    pub tick: u64,
    pub actor: [u8; 32],
    /// Signed look/move deltas (model-defined units).
    pub dx: i32,
    pub dy: i32,
    /// Button/key bitmask.
    pub buttons: u32,
    /// Free payload (e.g. a re-prompt), model-defined.
    pub payload: Vec<u8>,
}

/// Pack an action into a `CausalOp.mutation` body:
/// `dx(i32 LE) dy(i32 LE) buttons(u32 LE) payload_len(u32 LE) payload…`.
pub fn encode_action(dx: i32, dy: i32, buttons: u32, payload: &[u8]) -> Vec<u8> {
    let mut o = Vec::with_capacity(16 + payload.len());
    o.extend_from_slice(&dx.to_le_bytes());
    o.extend_from_slice(&dy.to_le_bytes());
    o.extend_from_slice(&buttons.to_le_bytes());
    o.extend_from_slice(&(payload.len() as u32).to_le_bytes());
    o.extend_from_slice(payload);
    o
}

/// Build and sign one action op (the canonical control event).
pub fn sign_action(
    key: &SigningKey, lamport: u64, tick: u64, dx: i32, dy: i32, buttons: u32, payload: &[u8],
) -> CausalOp {
    // Actions target the agent slot 0 (the world's controlled entity); the model
    // interprets the body. deps empty — actions are totally ordered by (tick,
    // lamport, actor), which is all replay-equivalence needs.
    CausalOp::sign(key, lamport, tick, vec![], [0u8; 16], MUT_ACTION, encode_action(dx, dy, buttons, payload))
}

/// Decode an action op's body into an [`Action`]. Errors if it is not an action op
/// or the body is malformed.
pub fn decode_action(op: &CausalOp) -> Result<Action, &'static str> {
    if op.mutation_kind != MUT_ACTION {
        return Err("not an action op");
    }
    let m = &op.mutation;
    if m.len() < 16 {
        return Err("action body truncated");
    }
    let dx = i32::from_le_bytes(m[0..4].try_into().unwrap());
    let dy = i32::from_le_bytes(m[4..8].try_into().unwrap());
    let buttons = u32::from_le_bytes(m[8..12].try_into().unwrap());
    let plen = u32::from_le_bytes(m[12..16].try_into().unwrap()) as usize;
    if m.len() != 16 + plen {
        return Err("action payload length mismatch");
    }
    Ok(Action { tick: op.effect_tick, actor: op.actor, dx, dy, buttons, payload: m[16..].to_vec() })
}

/// A replayed session: the portable session hash, the canonically-ordered action
/// sequence the model is fed, and the step count.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ModelSession {
    /// BLAKE3 hash-chain over the header + canonically-ordered action ops — the
    /// portable, replay-equivalent session identity. Equal on every sink that holds
    /// the same op set, regardless of arrival order.
    pub session_hash: [u8; 32],
    pub actions: Vec<Action>,
    pub n_steps: usize,
}

/// Replay an action trace → [`ModelSession`]. Canonically orders the ops
/// (`order_key` = `(effect_tick, lamport, actor)`), optionally verifies every
/// signature, and hash-chains `header ‖ op₀ ‖ … ‖ opₙ` under
/// [`DOMAIN_MODEL_SESSION`]. The result is identical on every machine — the
/// **replay-equivalence** the float renderer can't itself provide.
pub fn replay(header: &ModelHeader, ops: &[CausalOp], verify_sigs: bool) -> Result<ModelSession, &'static str> {
    let mut ordered: Vec<&CausalOp> = ops.iter().collect();
    ordered.sort_by_key(|o| o.order_key());

    let mut chain = {
        let mut h = blake3::Hasher::new();
        h.update(DOMAIN_MODEL_SESSION);
        h.update(&header.preimage());
        *h.finalize().as_bytes()
    };
    let mut actions = Vec::with_capacity(ordered.len());
    for op in ordered {
        if op.mutation_kind != MUT_ACTION {
            return Err("non-action op in a world-model trace");
        }
        if verify_sigs && op.verify().is_err() {
            return Err("action signature invalid");
        }
        actions.push(decode_action(op)?);
        let mut h = blake3::Hasher::new();
        h.update(DOMAIN_MODEL_SESSION);
        h.update(&chain);
        h.update(&op.to_bytes());
        chain = *h.finalize().as_bytes();
    }
    Ok(ModelSession { session_hash: chain, actions, n_steps: ops.len() })
}

/// Build a signed, energy-accounted **provenance receipt** for the frames a world
/// model generated from this session. Each frame is a content-addressed lineage step
/// carrying its measured microjoules; the chain originates at the `session_hash` (the
/// trace the frames were rendered from) and is Merkle-bound + Ed25519-sealed by
/// [`AssetReceipt`]. When the renderer is a float model the frames aren't
/// reproducible, so this receipt is the trust: who rendered, from which session, at
/// what energy. `frames` is `(frame content hash, microjoules)`, in display order;
/// must be non-empty. The receipt's `asset_hash` is the last frame.
pub fn frame_receipt(
    signer: &SigningKey, signer_id: impl Into<String>, session_hash: [u8; 32], frames: &[([u8; 32], u64)],
) -> AssetReceipt {
    let mut steps = Vec::with_capacity(frames.len());
    let mut input = session_hash; // the session is the originating input
    for (i, (fh, j)) in frames.iter().enumerate() {
        steps.push(ProvStep { op: format!("render.frame.{i}"), input_hash: input, output_hash: *fh, work: i as u64, joules_micro: *j });
        input = *fh;
    }
    let asset = frames.last().map(|f| f.0).unwrap_or(session_hash);
    let group_id = u64::from_le_bytes(session_hash[..8].try_into().unwrap());
    AssetReceipt::build(signer, signer_id, asset, "video", "wai.world.model", group_id, steps, None)
}

/// Lower-case hex of a 32-byte hash (session hashes ride in JSON/receipts as hex).
pub fn hex32(b: &[u8; 32]) -> String {
    let mut s = String::with_capacity(64);
    for &x in b {
        s.push(char::from_digit((x >> 4) as u32, 16).unwrap());
        s.push(char::from_digit((x & 15) as u32, 16).unwrap());
    }
    s
}

/// Convenience: the `OpId` of an action (actor + lamport) — for building `deps` if a
/// model ever needs causal action ordering beyond the total order.
pub fn action_id(op: &CausalOp) -> OpId {
    op.id()
}