open-standards/joulehook/joulehook-rs/crates/joulehook-meter/src/lib.rs

//! JouleHook — the meter-seam binding (`joulehook/1`, SPEC §4 Proxy).
//!
//! The gate seals every *tool* call as a metered transaction; the meter does the
//! same for the *model* call. Two composable primitives a proxy (or a host's
//! request/response hooks) wires around the model endpoint:
//!
//! - [`inject_ceiling`] — write a JouleHook grant's remaining joule budget into the
//!   request's `metadata.joule_ceiling_uj`, never *raising* a tighter ceiling. This
//!   is what turns *metered* into *capped*: over budget, Mesh refuses the request
//!   before it reaches the provider.
//! - [`bridge_receipt`] — read Mesh's response `usage` (tokens + `usage.energy`) and
//!   seal it as an `efficiency/1` receipt, carrying the honest provenance straight
//!   through from `usage.energy.provenance`.
//!
//! The authoritative energy figure comes from Mesh's meter, not an estimate — this
//! binding only carries it into the same receipt surface the gate uses.

use ed25519_dalek::SigningKey;
use jcp_efficiency::Efficiency;
use jcp_schema::{ContentId, Energy};
use serde_json::Value;

/// Write `budget_uj` into `request.metadata.joule_ceiling_uj`, clamped so an
/// existing *tighter* ceiling is preserved (`min`). The request is returned with the
/// ceiling set; a non-object request is returned unchanged.
pub fn inject_ceiling(mut request: Value, budget_uj: u64) -> Value {
    if let Some(obj) = request.as_object_mut() {
        let metadata = obj
            .entry("metadata")
            .or_insert_with(|| Value::Object(Default::default()));
        if let Some(m) = metadata.as_object_mut() {
            let ceiling = m
                .get("joule_ceiling_uj")
                .and_then(Value::as_u64)
                .map_or(budget_uj, |existing| existing.min(budget_uj));
            m.insert("joule_ceiling_uj".into(), Value::from(ceiling));
        }
    }
    request
}

/// The reference meter steward (a fixed dev key). A production deployment supplies
/// its own steward key + DID.
pub const METER_STEWARD: &str = "did:jcp:joulehook-meter";

/// The reference meter steward signing key.
pub fn meter_steward_sk() -> SigningKey {
    SigningKey::from_bytes(&[0x6e; 32])
}

/// [`bridge_receipt`] under the reference meter steward, seq 0 / genesis prev — the
/// convenience the proxy and CLI use so callers don't wrangle keys.
pub fn bridge_receipt_default(
    response: &Value,
    task: &str,
    completed: bool,
    verdict: &str,
    wall_ns: u64,
) -> Option<Vec<u8>> {
    bridge_receipt(
        response,
        task,
        completed,
        verdict,
        wall_ns,
        &meter_steward_sk(),
        METER_STEWARD,
        0,
        &ContentId([0u8; 32]),
    )
}

/// Map Mesh's honest provenance label onto a [`jcp_schema::Energy`]. Mesh writes
/// the family's wire labels (`HwShunt`, `ModelBased`, `Estimator` — its
/// `Provenance::jcr1_label`); the lower-case display spelling (`hw_shunt`,
/// `model_based`, `estimator`) that JouleClaw's energy crate prints is read as the
/// same grade. Any other label falls to the estimator floor rather than
/// overclaiming a measurement.
fn energy_from(uj: u64, provenance: &str) -> Energy {
    match provenance {
        "HwShunt" | "hw_shunt" => Energy::measured(uj),
        "ModelBased" | "model_based" => Energy::modeled(uj),
        _ => Energy::estimator(uj),
    }
}

/// Seal a completed model serve as an `efficiency/1` receipt from Mesh's response
/// `usage`. Returns `None` if the response carries no `usage.energy` block (i.e. it
/// did not come through a Mesh meter — this binding does not invent an energy
/// figure). `completed`/`verdict` are the caller's outcome anchor for the serve.
#[allow(clippy::too_many_arguments)]
pub fn bridge_receipt(
    response: &Value,
    task: &str,
    completed: bool,
    verdict: &str,
    wall_ns: u64,
    steward_sk: &SigningKey,
    steward: &str,
    seq: u64,
    prev: &ContentId,
) -> Option<Vec<u8>> {
    let usage = response.get("usage")?;
    let energy = usage.get("energy")?;
    let uj = energy.get("uj").and_then(Value::as_u64)?;
    let provenance = energy
        .get("provenance")
        .and_then(Value::as_str)
        .unwrap_or("estimator");
    let eff = Efficiency {
        task: task.to_string(),
        completed,
        verdict: verdict.to_string(),
        tokens_in: usage.get("prompt_tokens").and_then(Value::as_u64).unwrap_or(0),
        tokens_out: usage
            .get("completion_tokens")
            .and_then(Value::as_u64)
            .unwrap_or(0),
        energy: energy_from(uj, provenance),
        wall_ns,
        value_uusd: None,
        energy_attestation: None,
    };
    Some(eff.seal(seq, prev, steward_sk, steward))
}

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

    #[test]
    fn injects_the_ceiling_when_absent() {
        let req = serde_json::json!({ "model": "m", "messages": [] });
        let out = inject_ceiling(req, 5_000_000);
        assert_eq!(out["metadata"]["joule_ceiling_uj"], 5_000_000);
    }

    #[test]
    fn never_raises_a_tighter_existing_ceiling() {
        // A caller asking for ≤ 1 mJ must not be loosened to the grant's 5 mJ.
        let req = serde_json::json!({ "metadata": { "joule_ceiling_uj": 1_000_000 } });
        let out = inject_ceiling(req, 5_000_000);
        assert_eq!(out["metadata"]["joule_ceiling_uj"], 1_000_000);
    }

    #[test]
    fn lowers_a_looser_existing_ceiling_to_the_budget() {
        let req = serde_json::json!({ "metadata": { "joule_ceiling_uj": 9_000_000 } });
        let out = inject_ceiling(req, 5_000_000);
        assert_eq!(out["metadata"]["joule_ceiling_uj"], 5_000_000);
    }

    #[test]
    fn bridges_mesh_usage_into_a_verifiable_efficiency_receipt() {
        let sk = SigningKey::from_bytes(&[0x6e; 32]);
        let steward = "did:jcp:joulehook-meter";
        let response = serde_json::json!({
            "usage": {
                "prompt_tokens": 8000,
                "completion_tokens": 900,
                "total_tokens": 8900,
                "energy": { "uj": 2_700_000, "uj_per_accepted_token": 3000, "provenance": "ModelBased" }
            }
        });
        let sealed = bridge_receipt(
            &response, "fix-auth-bug", true, "tests-pass", 42_000_000_000,
            &sk, steward, 0, &ContentId([0u8; 32]),
        )
        .expect("a usage.energy block seals");

        let pk = sk.verifying_key().to_bytes();
        let r = jcp_efficiency::verify(&sealed, steward, &pk).expect("efficiency/1 verifies");
        assert_eq!(r.tokens_out, 900);
        assert_eq!(r.energy_uj, 2_700_000);
        assert_eq!(r.uj_per_accepted_token, Some(2_700_000 / 900)); // 3000 µJ/accepted token
        assert_eq!(r.provenance, "ModelBased");
    }

    /// Bridge a real Mesh response — `mesh-openai`'s chat completion, its
    /// `usage.energy.provenance` written by Mesh's own wire label — and return
    /// the provenance the sealed receipt carries.
    fn bridged_grade_of_mesh_output(p: mesh_schema::Provenance) -> String {
        let response = mesh_openai::chat_completion_response(
            "chatcmpl-1", "m", 1_700_000_000, "ok", 12, 3, 66_900_000, Some(22_300_000),
            p.jcr1_label(),
        );
        let sealed = bridge_receipt_default(&response, "t", true, "ok", 0).expect("a usage.energy block seals");
        let pk = meter_steward_sk().verifying_key().to_bytes();
        jcp_efficiency::verify(&sealed, METER_STEWARD, &pk).expect("efficiency/1 verifies").provenance
    }

    /// Mesh writes its grades as `HwShunt` / `ModelBased` / `Estimator`. Each must
    /// reach the receipt as itself: a Mesh figure is neither under-graded to the
    /// estimator floor nor raised.
    #[test]
    fn mesh_grades_reach_the_receipt_unchanged() {
        use mesh_schema::Provenance as Mesh;
        for p in [Mesh::HwShunt, Mesh::ModelBased, Mesh::Estimator] {
            assert_eq!(bridged_grade_of_mesh_output(p), p.jcr1_label(), "Mesh's {p:?} figure");
        }
    }

    /// The lower-case spelling (`hw_shunt` / `model_based` / `estimator`, the
    /// JouleClaw energy crate's display form) is still read as the same grade, and
    /// a label the family does not write falls to the estimator floor.
    #[test]
    fn lower_case_labels_and_unknown_labels() {
        for (label, grade) in [
            ("hw_shunt", "HwShunt"),
            ("model_based", "ModelBased"),
            ("estimator", "Estimator"),
            ("measured", "Estimator"),
            ("hwshunt", "Estimator"),
            ("", "Estimator"),
        ] {
            let response = mesh_openai::chat_completion_response(
                "chatcmpl-1", "m", 1_700_000_000, "ok", 12, 3, 66_900_000, None, label,
            );
            let sealed = bridge_receipt_default(&response, "t", true, "ok", 0).expect("seals");
            let pk = meter_steward_sk().verifying_key().to_bytes();
            let r = jcp_efficiency::verify(&sealed, METER_STEWARD, &pk).expect("verifies");
            assert_eq!(r.provenance, grade, "label {label:?}");
        }
    }

    #[test]
    fn no_energy_block_seals_nothing() {
        let sk = SigningKey::from_bytes(&[0x6e; 32]);
        let response = serde_json::json!({ "usage": { "prompt_tokens": 1, "completion_tokens": 1 } });
        assert!(bridge_receipt(
            &response, "t", true, "ok", 0, &sk, "did:jcp:joulehook-meter", 0, &ContentId([0u8; 32])
        )
        .is_none());
    }
}