open-standards/wai/wai-rs/tests/jpegai_decode_path.rs

//! The reference dispatch reports the decode path that ran (extensions/determinism-tiers.md
//! §5), a receipt of the decode states it (extensions/jwp-receipts.md §2.4.12, a decoder
//! entry's `variant`), and the reference check of a reconstruction binding reads it
//! (extensions/reconstruction-binding.md §2).
//!
//! - A capability with one decode path reports `DecodePath::Sole`, and the capability it
//!   decoded through: the manifest's `capability`, or its `fallback`. A binding over a
//!   `decode-equivalence` capability's picture is checked: it holds, a changed hash
//!   differs, and a descriptor naming another picture is another reconstruction, never a
//!   failure. Committed inputs only; neither test skips.
//! - `wai.image.jpegai` without the integer weight bundle runs its float decode through
//!   ONNX Runtime. The result says `DecodePath::Float`; a session claim recorded from it
//!   says `"variant":"float"` and verifies, and the same decode named under a
//!   `NeuralIntegerExact` capability is refused (`session_variant_not_exact`); a
//!   reconstruction binding over its own pixels is reported as not checked, never as
//!   holding. Needs the float path's model bundle (`WAI_JPEGAI_ASSETS`, the local
//!   JPEG-AI asset directory exported from the ISO/IEC 6048-1 reference software; not in
//!   the repository). That directory is in no CI job, so this test runs locally only, on
//!   a host that holds it, as `tests/jpegai_live.rs` does (wai-jpegai-e2e.yml).
//! - With the integer weight bundle the exact decoder runs. The result says
//!   `DecodePath::Exact`, the claim `"variant":"exact"`, and the binding check reports
//!   whether the hash holds, and a binding of another `decoderId` as another
//!   reconstruction. Needs the bundle (`WAI_JPEGAI_WEIGHTS_DIR`, `WAI_JPEGAI_NORM`, as
//!   `tests/jpegai_e2e_full.rs` reads them), which `tools/wai_jpegai_fetch_e2e_inputs.sh`
//!   fetches; wai-neural-gates.yml runs it so, with `WAI_REQUIRE_SOTA_VECTORS=1`.
//!
//! Without their inputs the two JPEG-AI tests skip loudly, and with
//! `WAI_REQUIRE_SOTA_VECTORS=1` they fail (support/require_vectors.rs).
#![cfg(all(feature = "neural", feature = "receipt_claims", feature = "video_int"))]

#[path = "support/require_vectors.rs"]
mod require_vectors;

use std::path::{Path, PathBuf};

use ed25519_dalek::SigningKey;
use wai::claim_receipt::{verify_claim_json, ClaimBody, Rejection, SessionHost, SessionPath};
use wai::codecs::{CAP_IMAGE_JPEGAI, CAP_NEURAL_INT_HYPER, CAP_VIDEO_INT_HYPER, CAP_VIDEO_INT_MOTION};
use wai::container::{Conditioning, Manifest, ModelRequirement, Wai};
use wai::energy_report::{BuildProfile, BuildRecord};
use wai::neural::{check_reconstruction_binding, decode_envelope, BindingCheck, DecodePath, Decoded, Dispatched, ModelRegistry};
use wai::reconstruction_binding::{ReconstructionBinding, ReconstructionDescriptor};
use wai::session_recorder::{DecoderKey, SessionRecorder};

const TEST256: &[u8] = include_bytes!("vectors/jpegai/test256.bin");
const T: &str = "2026-11-04T12:00:00Z";

fn envelope(media: &str, req: ModelRequirement, payload: Vec<u8>) -> Wai {
    Wai::new(Manifest::new(media, "improve", req, Conditioning::new("none")), payload)
}

fn picture() -> Wai {
    envelope("image", ModelRequirement::new(CAP_IMAGE_JPEGAI), TEST256.to_vec())
}

/// The session claim a sink records of one decode of `env`, its decoder entry naming
/// `capability` on the path the dispatch reported, sealed and checked as
/// `SessionRecorder::finish` checks it.
fn record(env: &Wai, capability: &str, d: &Dispatched) -> Result<String, Rejection> {
    let Decoded::Image(img) = &d.output else { panic!("{}: not a picture", d.capability) };
    let bytes = env.to_bytes().expect("an envelope");
    let host = SessionHost { profile: "native-ffi".into(), target: None, build: BuildRecord { profile: BuildProfile::Release, toolchain: None, target: None }, threads: 1 };
    let mut rec = SessionRecorder::new(wai::pin::sha256(b"jpegai_decode_path: a session nonce"), "image", host, None);
    let entry = rec.decoder(DecoderKey {
        capability: capability.to_string(),
        rendition: None,
        object: wai::jwp_hash::content_hash(&bytes),
        first_unit: 0,
        path: SessionPath::Primary,
        decoder: format!("builtin:wai@{}", env!("CARGO_PKG_VERSION")),
        variant: d.path.variant().map(str::to_owned),
        pin: None,
        tier: None,
        module_status: None,
    });
    rec.received_envelope(&bytes, SessionPath::Primary);
    rec.presented(0, 1, entry, wai::pin::sha256(&img.rgb), None).expect("one presentation");
    let key = SigningKey::from_bytes(&[0x44; 32]);
    rec.finish(wai::jwp_hash::content_hash(&bytes), &key, "did:example:sink", 0, None, T).map(|s| s.claim.to_json())
}

/// The variant the one decoder entry of a verified session claim states.
fn variant_of(wire: &str) -> Option<String> {
    let c = verify_claim_json(wire).expect("the claim verifies");
    let ClaimBody::ReconstructSession(s) = &c.body else { panic!("a session claim") };
    s.decoders[0].variant.clone()
}

fn descriptor() -> ReconstructionDescriptor {
    ReconstructionDescriptor {
        decoder_id: 2,
        profile: "tools_0".into(),
        level: "L1".into(),
        width: 256,
        height: 256,
        format: "RGB8".into(),
        color_space: "bt709".into(),
    }
}

/// A capability with one decode path reports it, with the capability the payload decoded
/// through: the manifest's own, and the fallback when the sink does not hold the first.
#[test]
fn a_capability_with_one_path_reports_it_and_the_capability_it_decoded_through() {
    let clip = std::fs::read(Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_video/int_video_bitstream.wiv1")).unwrap();
    let reg = ModelRegistry::new().builtin(CAP_VIDEO_INT_MOTION);
    let own = decode_envelope(&envelope("video", ModelRequirement::new(CAP_VIDEO_INT_MOTION), clip.clone()), &reg).expect("decodes");
    assert_eq!((own.capability.as_str(), own.path, own.path.variant()), (CAP_VIDEO_INT_MOTION, DecodePath::Sole, None));
    let fallback = envelope("video", ModelRequirement::new(CAP_VIDEO_INT_HYPER).with_fallback(CAP_VIDEO_INT_MOTION), clip);
    let via = decode_envelope(&fallback, &reg).expect("decodes through the fallback");
    assert_eq!((via.capability.as_str(), via.path), (CAP_VIDEO_INT_MOTION, DecodePath::Sole));
    assert!(matches!(via.output, Decoded::Int(_)));
}

/// A binding over the picture of a `decode-equivalence` capability with one decode path
/// (`wai.neural.int_mlicpp`, the committed vector): the check decodes and compares, and a
/// descriptor that names another picture than the one decoded is reported as another
/// reconstruction, not as a failure.
#[test]
fn a_binding_over_a_decode_equivalence_capability_is_checked_under_its_descriptor() {
    let d = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_mlicpp");
    let env = Wai::from_bytes(&std::fs::read(d.join("stream.wai")).unwrap()).unwrap();
    let reg = ModelRegistry::new().register("wai.neural.int_mlicpp", d.join("model.wmb"));
    let out = decode_envelope(&env, &reg).expect("decodes");
    assert_eq!((out.capability.as_str(), out.path), ("wai.neural.int_mlicpp", DecodePath::Sole));
    let Decoded::Image(img) = &out.output else { panic!("not a picture") };
    let at = |width, height, format: &str| ReconstructionDescriptor {
        decoder_id: 0,
        profile: "int_mlicpp".into(),
        level: "-".into(),
        width,
        height,
        format: format.into(),
        color_space: "bt709".into(),
    };
    let binding = ReconstructionBinding::bind(&img.rgb, at(img.width, img.height, "RGB8"));
    assert_eq!(check_reconstruction_binding(&env, &reg, &binding).unwrap(), BindingCheck::Holds);
    let mut changed = binding.clone();
    changed.hash[0] ^= 1;
    assert_eq!(check_reconstruction_binding(&env, &reg, &changed).unwrap(), BindingCheck::Differs);
    for other in [at(img.width / 2, img.height * 2, "RGB8"), at(img.width, img.height, "RGBA8")] {
        let b = ReconstructionBinding { descriptor: other, ..binding.clone() };
        assert_eq!(check_reconstruction_binding(&env, &reg, &b).unwrap(), BindingCheck::OtherReconstruction(DecodePath::Sole));
    }
}

/// The float decode of `wai.image.jpegai` says so in the result, in the claim and in the
/// binding check.
#[test]
fn the_jpegai_float_fallback_reports_itself() {
    const TEST: &str = "jpegai_decode_path::the_jpegai_float_fallback_reports_itself";
    const PRODUCE: &str = "the local JPEG-AI asset directory wai/jpegai-vm-assets (exported from the \
ISO/IEC 6048-1 reference software; not in the repository), or WAI_JPEGAI_ASSETS=<dir>";
    let assets = require_vectors::env_dir(TEST, "WAI_JPEGAI_ASSETS", "../jpegai-vm-assets", PRODUCE);
    if !assets.join("model16/CCS_SGMM").exists() || !assets.join("jpegai_scaler_vec.json").exists() {
        assets.absent(TEST, &format!("the float path's model bundle is not in {}", assets.display()), PRODUCE);
        return;
    }
    assert!(!assets.join("norm.json").exists(), "the float path's bundle must not hold the exact decoder's");
    let env = picture();
    let reg = ModelRegistry::new().register(CAP_IMAGE_JPEGAI, assets.path());
    let d = decode_envelope(&env, &reg).expect("the float decode");
    assert_eq!((d.capability.as_str(), d.path, d.path.variant()), (CAP_IMAGE_JPEGAI, DecodePath::Float, Some("float")));
    let Decoded::Image(img) = &d.output else { panic!("not a picture") };
    assert_eq!((img.width, img.height, img.rgb.len()), (256, 256, 256 * 256 * 3));

    // The receipt states the path that ran, and verifies.
    let wire = record(&env, CAP_IMAGE_JPEGAI, &d).expect("a float decode of wai.image.jpegai is recorded under it");
    assert!(wire.contains("\"variant\":\"float\""), "{wire}");
    assert_eq!(variant_of(&wire).as_deref(), Some("float"));
    // Named under a NeuralIntegerExact capability, the same decode is no decode of it.
    let refused = record(&env, CAP_NEURAL_INT_HYPER, &d).expect_err("a float path under a NeuralIntegerExact capability");
    assert_eq!((refused.step, refused.reason.label()), (4, "session_variant_not_exact"));

    // A binding over this very decode's pixels is not checked by it, and not failed.
    let binding = ReconstructionBinding::bind(&img.rgb, descriptor());
    assert_eq!(check_reconstruction_binding(&env, &reg, &binding).unwrap(), BindingCheck::NotChecked(DecodePath::Float));
}

/// The exact decoder of `wai.image.jpegai` says so in the result and the claim, and the
/// binding check reads its pixels.
#[test]
fn the_jpegai_exact_decoder_reports_itself() {
    const TEST: &str = "jpegai_decode_path::the_jpegai_exact_decoder_reports_itself";
    const PRODUCE: &str = "tools/wai_jpegai_fetch_e2e_inputs.sh (the models root's jpegai/v1/, pinned by tools/jpegai-e2e-v1.sha256)";
    let Some([dir, norm]) = require_vectors::env_paths(TEST, ["WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE) else { return };
    // The bundle as the dispatch reads it: norm.json, and weights/<member>.json beside it.
    let root: PathBuf = std::env::temp_dir().join(format!("wai-jpegai-decode-path-{}", std::process::id()));
    std::fs::create_dir_all(root.join("weights")).unwrap();
    std::fs::copy(&norm, root.join("norm.json")).unwrap();
    for m in ["hs_y", "hs_uv", "hyp_y", "hyp_uv", "synth_y", "synth_uv", "mcm0_y", "mcm1_y", "mcm2_y", "mcm3_y", "icci"] {
        std::fs::copy(Path::new(&dir).join(format!("{m}.json")), root.join(format!("weights/{m}.json"))).unwrap();
    }
    let env = picture();
    let reg = ModelRegistry::new().register(CAP_IMAGE_JPEGAI, &root);
    let d = decode_envelope(&env, &reg).expect("the exact decode");
    assert_eq!((d.capability.as_str(), d.path, d.path.variant()), (CAP_IMAGE_JPEGAI, DecodePath::Exact, Some("exact")));
    let Decoded::Image(img) = &d.output else { panic!("not a picture") };

    let wire = record(&env, CAP_IMAGE_JPEGAI, &d).expect("an exact decode is recorded");
    assert_eq!(variant_of(&wire).as_deref(), Some("exact"));

    let binding = ReconstructionBinding::bind(&img.rgb, descriptor());
    assert_eq!(check_reconstruction_binding(&env, &reg, &binding).unwrap(), BindingCheck::Holds);
    let mut other = img.rgb.clone();
    other[0] ^= 1;
    let changed = ReconstructionBinding::bind(&other, descriptor());
    assert_eq!(check_reconstruction_binding(&env, &reg, &changed).unwrap(), BindingCheck::Differs);
    // A binding of another synthesis branch is over a reconstruction this decoder does not
    // compute: not checked, and not failed, though its hash is this decode's.
    let branch_1 = ReconstructionBinding::bind(&img.rgb, ReconstructionDescriptor { decoder_id: 1, ..descriptor() });
    assert_eq!(check_reconstruction_binding(&env, &reg, &branch_1).unwrap(), BindingCheck::OtherReconstruction(DecodePath::Exact));
    std::fs::remove_dir_all(&root).ok();
}