open-standards/wai/wai-quantum-service/tests/tasks.rs

//! A run as an MCP task, against the service itself: made at once, followed
//! by tasks/get to its end, and its result from tasks/result, its receipt
//! verifying; and tasks/result asked while the task still works, answered as
//! an event stream once it ends.

use serde_json::{Value, json};
use std::io::{Read, Write};
use std::net::{TcpListener, TcpStream};
use std::path::Path;
use std::process::{Child, Command, Stdio};
use std::time::{Duration, Instant};

fn free_port() -> u16 {
    TcpListener::bind("127.0.0.1:0").unwrap().local_addr().unwrap().port()
}

/// One HTTP/1.1 exchange: the status, the headers and the body.
fn http(port: u16, method: &str, path: &str, body: Option<&Value>, accept: &str) -> (u16, String, String) {
    let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap();
    s.set_read_timeout(Some(Duration::from_secs(180))).unwrap();
    let text = body.map(Value::to_string).unwrap_or_default();
    write!(s, "{method} {path} HTTP/1.1\r\nHost: 127.0.0.1\r\nContent-Type: application/json\r\nAccept: {accept}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{text}", text.len()).unwrap();
    let mut out = String::new();
    s.read_to_string(&mut out).unwrap();
    let status = out.split_whitespace().nth(1).and_then(|c| c.parse().ok()).unwrap_or(0);
    let (head, body) = out.split_once("\r\n\r\n").map(|(h, b)| (h.to_string(), b.to_string())).unwrap_or_default();
    (status, head, body)
}

fn rpc(port: u16, method: &str, params: Value) -> Value {
    let (code, _, body) = http(port, "POST", "/mcp", Some(&json!({ "jsonrpc": "2.0", "id": 1, "method": method, "params": params })), "application/json, text/event-stream");
    assert_eq!(code, 200, "{method}: {body}");
    serde_json::from_str(&body).unwrap_or_else(|e| panic!("{method}: {e}: {body}"))
}

fn start(dir: &Path, port: u16) -> Child {
    let child = Command::new(env!("CARGO_BIN_EXE_wai-quantum-service"))
        .arg("serve")
        .env("WAI_QSVC_ADDR", format!("127.0.0.1:{port}"))
        .env("WAI_QSVC_STORE", dir.join("store"))
        .env("WAI_QSVC_KEY", dir.join("signing.key"))
        .env("WAI_QSVC_JOB_GB", "2")
        .env("WAI_QSVC_BUDGET_GB", "4")
        .env("WAI_QSVC_WORKERS", "1")
        .env("WAI_QSVC_ENERGY", "0")
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .expect("the service starts");
    let t = Instant::now();
    while TcpStream::connect(("127.0.0.1", port)).is_err() || http(port, "GET", "/v1/health", None, "application/json").0 != 200 {
        assert!(t.elapsed() < Duration::from_secs(60), "the service did not answer");
        std::thread::sleep(Duration::from_millis(50));
    }
    child
}

#[test]
fn a_run_as_a_task_is_followed_to_its_end_and_its_result_streams_when_asked_early() {
    let dir = std::env::temp_dir().join(format!("wai-qsvc-tasks-{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&dir);
    std::fs::create_dir_all(&dir).unwrap();
    let port = free_port();
    let mut svc = start(&dir, port);

    // Made at once, followed by tasks/get to its end.
    let bell = json!({ "qasm": "OPENQASM 3.0; qubit[2] q; bit[2] c; h q[0]; cx q[0], q[1]; c = measure q;", "shots": 200, "seed": 3 });
    let made = rpc(port, "tools/call", json!({ "name": "run_quantum_circuit", "arguments": bell, "task": { "ttl": 60000 } }));
    let id = made["result"]["task"]["taskId"].as_str().unwrap_or_else(|| panic!("{made}")).to_string();
    assert_eq!(made["result"]["task"]["status"], "working", "{made}");
    let t = Instant::now();
    let ended = loop {
        let got = rpc(port, "tasks/get", json!({ "taskId": id }));
        if got["result"]["status"] != "working" {
            break got;
        }
        assert!(t.elapsed() < Duration::from_secs(120), "never ended: {got}");
        std::thread::sleep(Duration::from_millis(50));
    };
    assert_eq!(ended["result"]["status"], "completed", "{ended}");
    let r = rpc(port, "tasks/result", json!({ "taskId": id }));
    let sc = &r["result"]["structuredContent"];
    assert_eq!((sc["id"].as_str(), sc["status"].as_str(), r["result"]["isError"].as_bool()), (Some(id.as_str()), Some("done"), Some(false)), "{r}");
    assert_eq!(r["result"]["_meta"]["io.modelcontextprotocol/related-task"]["taskId"], id.as_str());
    assert!(wai_quantum::engine::job::verify(&json!({ "request": sc["request"], "result": sc["result"] })).unwrap().0, "{sc}");
    assert_eq!(sc["result"]["counts"].as_object().unwrap().values().filter_map(Value::as_u64).sum::<u64>(), 200);

    // Asked for its result while it works, the answer is an event stream
    // that ends with the result.
    // Asked a few milliseconds after it is made, it is still queued or running.
    let mut long = String::from("OPENQASM 3.0;\nqubit[16] q;\nh q;\n");
    for k in 0..1000 {
        long += &format!("rx(0.{}) q[{}];\ncx q[{}], q[{}];\n", k % 9 + 1, k % 16, k % 16, (k + 1) % 16);
    }
    let made = rpc(port, "tools/call", json!({ "name": "run_quantum_circuit", "arguments": { "qasm": long, "observables": ["Z0"], "method": "statevector" }, "task": {} }));
    let id = made["result"]["task"]["taskId"].as_str().unwrap_or_else(|| panic!("{made}")).to_string();
    let (code, head, body) = http(port, "POST", "/mcp", Some(&json!({ "jsonrpc": "2.0", "id": 9, "method": "tasks/result", "params": { "taskId": id } })), "application/json, text/event-stream");
    assert_eq!(code, 200, "{body}");
    assert!(head.to_ascii_lowercase().contains("content-type: text/event-stream"), "{head}");
    let data: String = body.lines().filter_map(|l| l.strip_prefix("data:")).map(str::trim_start).collect();
    let answer: Value = serde_json::from_str(&data).unwrap_or_else(|e| panic!("{e}: {body}"));
    assert_eq!((answer["id"].as_u64(), answer["result"]["structuredContent"]["status"].as_str()), (Some(9), Some("done")), "{answer}");
    assert_eq!(answer["result"]["_meta"]["io.modelcontextprotocol/related-task"]["taskId"], id.as_str());
    assert_eq!(rpc(port, "tasks/get", json!({ "taskId": id }))["result"]["status"], "completed");

    svc.kill().unwrap();
    svc.wait().unwrap();
    let _ = std::fs::remove_dir_all(&dir);
}