open-standards/mesh/mesh-rs/crates/mesh-openai/tests/serve_example.rs

//! The reference server example (`examples/serve.rs`) seals what its meter
//! read, never a fixture's figure.
//!
//! Driven here through one real `/v1/chat/completions` request over a socket,
//! with the example's own meter and mesh (the example file is compiled into
//! this test as a module, so what is tested is what `cargo run --example
//! serve` runs). Its receipt's grade must be one this build's meter can give:
//! with no hardware meter feature, only `Estimator`, with no measured figure.
//! A hardware figure must have been read: with `metal` (this process's own
//! counters) it must lie within this process's CPU energy read around the
//! request by a reader that shares no code with the meter — the kernel's
//! `ri_energy_nj`, read raw. A fixture that steps a fixed 4.8 J per snapshot
//! and calls it `HwShunt` fails both: its grade in a build without a
//! hardware meter, and the bound in a build with one (a CPU-only serve of 16
//! tokens is microjoules to millijoules here, nowhere near joules).
//!
//! And in every branch the request is billed something: a served request
//! ran, so a figure of 0 µJ — an estimator too coarse to see a span shorter
//! than its resolution — bills real work nothing.

#[allow(dead_code)]
#[path = "../examples/serve.rs"]
mod serve;

use std::io::{Read, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::sync::mpsc;
use std::thread;

use serde_json::Value;

#[cfg(target_os = "macos")]
unsafe extern "C" {
    fn proc_pid_rusage(pid: i32, flavor: i32, buffer: *mut std::ffi::c_void) -> i32;
}

/// This process's CPU energy, nJ: `ri_energy_nj` of `rusage_info_v6`, read
/// here into a raw buffer at the SDK's offset (336), not through any meter.
/// `None` off macOS, where no such per-process count exists.
fn cpu_nj() -> Option<u64> {
    #[cfg(target_os = "macos")]
    {
        let mut words = [0u64; 464 / 8];
        let rc = unsafe { proc_pid_rusage(std::process::id() as i32, 6, words.as_mut_ptr().cast()) };
        assert_eq!(rc, 0, "proc_pid_rusage");
        Some(words[336 / 8])
    }
    #[cfg(not(target_os = "macos"))]
    {
        None
    }
}

/// POST one chat completion and return the response's JSON body.
fn post(addr: SocketAddr) -> Value {
    let body = r#"{"model":"gpt-mini","messages":[{"role":"user","content":"hi"}],"max_tokens":16}"#;
    let mut conn = TcpStream::connect(addr).expect("connect");
    let req = format!(
        "POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
        body.len(),
        body
    );
    conn.write_all(req.as_bytes()).expect("send");
    let mut raw = String::new();
    conn.read_to_string(&mut raw).expect("read");
    assert!(raw.starts_with("HTTP/1.1 200 OK"), "status line: {}", &raw[..raw.len().min(80)]);
    serde_json::from_str(raw.split("\r\n\r\n").nth(1).expect("response body")).expect("JSON body")
}

#[test]
fn the_reference_server_seals_what_its_meter_read() {
    let listener = TcpListener::bind("127.0.0.1:0").expect("bind");
    let addr = listener.local_addr().expect("addr");
    let (built, meter_built) = mpsc::channel();
    // The server lives in one thread, as `main` runs it: the meter, the mesh
    // and its receipt chain are built and stay there.
    let server = thread::spawn(move || {
        let meter = serve::meter();
        built.send(meter.name().to_string()).expect("report the meter");
        let mut mesh = serve::mesh(meter.as_ref());
        mesh_openai::server::serve_n(&listener, &mut mesh, &serve::ByteBackend, 100_000_000, 1).expect("serve");
    });
    // Building the meter runs its liveness probes (a third of a second of CPU
    // and a GPU load, for `ProcessMeter`): wait for it, so the window below
    // holds the request and not the probes.
    let meter = meter_built.recv().expect("the server built its meter");

    let before = cpu_nj();
    let v = post(addr);
    let around_uj = cpu_nj().zip(before).map(|(after, before)| (after - before) / 1_000);
    server.join().expect("server thread");

    let energy = &v["usage"]["energy"];
    let receipt = &energy["receipt"];
    let provenance = energy["provenance"].as_str().expect("usage.energy.provenance");
    let measured = receipt.get("measured_uj").and_then(Value::as_u64);
    eprintln!(
        "serve example: meter {meter}; usage.energy {} µJ [{provenance}]; sealed {:?} measured {measured:?} estimated {}; \
         this process's CPU energy around the request {around_uj:?} µJ",
        energy["uj"], receipt["energy_provenance"], receipt["estimated_uj"]
    );

    assert_eq!(receipt["mesh_receipt"], "mesh/1");
    let uj = energy["uj"].as_u64().expect("usage.energy.uj");
    assert!(uj > 0, "a served request was billed 0 µJ [{provenance}] by meter {meter}: {energy}");
    assert_eq!(
        receipt["energy_provenance"].as_str(),
        Some(provenance),
        "the plaintext grade and the sealed one differ: {energy}"
    );
    match provenance {
        "Estimator" => assert_eq!(measured, None, "an estimate sealed with a measured figure: {receipt}"),
        "HwShunt" => {
            if !cfg!(feature = "hw-meter") {
                panic!("a build with no hardware meter sealed a hardware reading, which nothing read: {receipt}");
            }
            let uj = measured.expect("a HwShunt figure carries its measured figure");
            assert!(uj > 0, "a served span read no energy: {receipt}");
            // Only this process's own counters bound it by this process's
            // energy; a machine counter reads the whole host.
            if cfg!(feature = "metal") {
                let around = around_uj.expect("macOS reads this process's CPU energy");
                assert!(
                    uj <= around + 1,
                    "the receipt read {uj} µJ, more than this process's {around} µJ of CPU energy around a CPU-only serve: \
                     a figure nothing read"
                );
            }
        }
        "ModelBased" => {
            if !cfg!(feature = "hw-meter") || cfg!(feature = "metal") {
                panic!("no meter of this build grades a span by a model: {receipt}");
            }
        }
        other => panic!("an unknown grade `{other}`: {receipt}"),
    }
    if !cfg!(feature = "hw-meter") {
        assert_eq!(provenance, "Estimator", "a build with no hardware meter can only estimate: {receipt}");
    }
}