open-standards/jcp/jcp-rs/crates/jcp-gateway/tests/stdio_binary.rs

//! The `jcp-gateway` binary itself, over stdio, in front of a real server
//! process: the way an MCP client runs it. Each test drives the binary with
//! newline-delimited JSON-RPC on its stdin and reads its stdout on a thread
//! of its own, so a gateway that stops answering fails the test at a
//! deadline instead of hanging it.
//!
//! The server is awk, answering only the lines that carry an `id`, as an MCP
//! server answers requests and never notifications. It runs on every host
//! this crate builds on, whichever awk the host has ([`awk_server`]).

use serde_json::{json, Value};
use std::io::{BufRead, BufReader, Read, Write};
use std::process::{Child, ChildStdin, Command, Stdio};
use std::sync::mpsc::{self, Receiver};
use std::time::Duration;

/// How long a test waits for one line from the gateway: far past anything a
/// call takes, on a host at a load of 14 to 260.
const LINE_DEADLINE: Duration = Duration::from_secs(120);

/// The command line of an awk server running `script` over a pipe.
///
/// An awk reading a pipe must hand each line to the script as it arrives, or
/// the server never answers a request until its client closes the pipe. mawk,
/// the default awk on Debian and Ubuntu hosts, buffers piped input in blocks
/// unless it is run with `-W interactive`; the other awks read lines as they
/// come and reject or ignore that option. So the awk is asked what it is: a
/// host whose awk is mawk once failed both tests at their deadline, while the
/// same tests passed where the awk was another.
fn awk_server(script: &str) -> Vec<String> {
    const AWK: &str = "/usr/bin/awk";
    let is_mawk = Command::new(AWK)
        .args(["-W", "version"])
        .stdin(Stdio::null())
        .output()
        .map(|o| {
            String::from_utf8_lossy(&o.stdout).contains("mawk") || String::from_utf8_lossy(&o.stderr).contains("mawk")
        })
        .unwrap_or(false);
    let mut argv = vec![AWK.to_string()];
    if is_mawk {
        argv.extend(["-W".to_string(), "interactive".to_string()]);
    }
    argv.push(script.to_string());
    argv
}

/// An MCP-ish server in awk: a request (a line with an `id`) gets one
/// response carrying that id; a line without one gets nothing.
const ANSWERS_REQUESTS_ONLY: &str = r#"/"id"/ {
  id = "null"
  if (match($0, /"id":[0-9]+/)) id = substr($0, RSTART + 5, RLENGTH - 5)
  if ($0 ~ /"method":"initialize"/) printf "{\"jsonrpc\":\"2.0\",\"id\":%s,\"result\":{\"protocolVersion\":\"2025-06-18\",\"capabilities\":{}}}\n", id
  else printf "{\"jsonrpc\":\"2.0\",\"id\":%s,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"done\"}]}}\n", id
  fflush()
}"#;

/// Where the gateway run under `name` keeps its receipt stream.
fn receipts_path(name: &str) -> std::path::PathBuf {
    std::path::Path::new(env!("CARGO_TARGET_TMPDIR")).join(format!("jcp-gateway-{name}-{}.cose", std::process::id()))
}

/// The gateway binary, running, with its server.
struct Gateway {
    child: Child,
    stdin: ChildStdin,
    lines: Receiver<String>,
    stderr: Receiver<String>,
}

impl Gateway {
    fn spawn(name: &str, server: &[&str], extra: &[&str]) -> Gateway {
        let _ = std::fs::remove_file(receipts_path(name));
        Gateway::restart(name, server, extra)
    }

    /// The gateway over the receipt stream an earlier run under `name` left,
    /// as a restart finds it.
    fn restart(name: &str, server: &[&str], extra: &[&str]) -> Gateway {
        let receipts = receipts_path(name);
        let mut child = Command::new(env!("CARGO_BIN_EXE_jcp-gateway"))
            .args(["--server", name, "--grant", "echo=1000000000"])
            .args(extra)
            .arg("--")
            .args(server)
            .env("JCP_GATEWAY_RECEIPTS", &receipts)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .spawn()
            .expect("spawn the jcp-gateway binary");
        let stdin = child.stdin.take().unwrap();
        let (tx, lines) = mpsc::channel();
        let out = child.stdout.take().unwrap();
        std::thread::spawn(move || {
            for line in BufReader::new(out).lines() {
                let Ok(line) = line else { break };
                if tx.send(line).is_err() {
                    break;
                }
            }
        });
        let (etx, stderr) = mpsc::channel();
        let mut err = child.stderr.take().unwrap();
        std::thread::spawn(move || {
            let mut all = String::new();
            let _ = err.read_to_string(&mut all);
            let _ = etx.send(all);
        });
        Gateway { child, stdin, lines, stderr }
    }

    fn send(&mut self, msg: &Value) {
        writeln!(self.stdin, "{msg}").expect("write to the gateway");
        self.stdin.flush().expect("flush to the gateway");
    }

    /// The next line the gateway writes, parsed; fails the test, with the
    /// gateway's stderr, if none comes by the deadline.
    fn next(&mut self, waiting_for: &str) -> Value {
        match self.lines.recv_timeout(LINE_DEADLINE) {
            Ok(line) => serde_json::from_str(&line).unwrap_or_else(|e| panic!("not JSON ({e}): {line}")),
            Err(_) => {
                let _ = self.child.kill();
                let _ = self.child.wait();
                let stderr = self.stderr.recv_timeout(Duration::from_secs(5)).unwrap_or_default();
                panic!("no line from the gateway while waiting for {waiting_for}; its stderr:\n{stderr}")
            }
        }
    }
}

impl Drop for Gateway {
    fn drop(&mut self) {
        // Its server reads the gateway's pipe, so it sees EOF and exits too.
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

fn call(id: i64) -> Value {
    json!({"jsonrpc":"2.0","id":id,"method":"tools/call","params":{"name":"echo","arguments":{}}})
}

/// The handshake every MCP client opens with — `initialize`, then the
/// `notifications/initialized` notification, which the server answers with
/// nothing — and then calls. The gateway must send the notification on
/// without waiting for a reply, or it never reads another line.
#[test]
fn the_binary_serves_calls_after_the_initialized_notification() {
    let mut g = Gateway::spawn("smoke", &awk_server(ANSWERS_REQUESTS_ONLY).iter().map(String::as_str).collect::<Vec<_>>(), &[]);
    g.send(&json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{}}}));
    let init = g.next("the initialize response");
    assert_eq!(init["id"], json!(1), "{init}");
    assert!(init["result"]["capabilities"]["extensions"]["science.transaction/jcp"].is_object(), "{init}");
    g.send(&json!({"jsonrpc":"2.0","method":"notifications/initialized"}));
    for id in 2..5 {
        g.send(&call(id));
        let r = g.next(&format!("the response to call {id}"));
        assert_eq!(r["id"], json!(id), "{r}");
        assert_eq!(r["result"]["content"][0]["text"], json!("done"), "{r}");
        let meta = &r["result"]["_meta"]["science.transaction/jcp"];
        assert!(meta["receipt_b64"].is_string(), "each call carries its sealed receipt: {r}");
        eprintln!("stdio binary: call {id}: cost_uj {} [{}] domain {}", meta["cost_uj"], meta["energy_provenance"], meta["energy_domain"]);
        // A notification between calls is sent on, unanswered, too.
        g.send(&json!({"jsonrpc":"2.0","method":"notifications/cancelled","params":{"requestId":99}}));
    }
}

/// An MCP-ish server in awk that, inside each `tools/call`, first sends the
/// client a `notifications/progress` and a request of its own
/// (`sampling/createMessage`, id `s<call id>`), waits for the client's
/// answer, and only then sends the call's response, saying whether it got
/// the answer. Every other request gets a plain response; notifications get
/// nothing.
const TALKS_MID_CALL: &str = r#"/"method":"tools\/call"/ {
  id = "null"
  if (match($0, /"id":[0-9]+/)) id = substr($0, RSTART + 5, RLENGTH - 5)
  printf "{\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progressToken\":%s,\"progress\":1}}\n", id
  printf "{\"jsonrpc\":\"2.0\",\"id\":\"s%s\",\"method\":\"sampling/createMessage\",\"params\":{\"maxTokens\":1}}\n", id
  fflush()
  answer = ""
  getline answer
  got = (answer ~ ("\"id\":\"s" id "\"") && answer ~ /"result"/) ? "sampled" : "unanswered"
  printf "{\"jsonrpc\":\"2.0\",\"id\":%s,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"%s\"}]}}\n", id, got
  fflush()
  next
}
/"id"/ {
  id = "null"
  if (match($0, /"id":[0-9]+/)) id = substr($0, RSTART + 5, RLENGTH - 5)
  printf "{\"jsonrpc\":\"2.0\",\"id\":%s,\"result\":{\"protocolVersion\":\"2025-06-18\",\"capabilities\":{}}}\n", id
  fflush()
}"#;

/// A server that talks in the middle of a call: the gateway hands its
/// notification and its request to the client as they come, carries the
/// client's answer back to the server while the call is still open, and
/// returns, as the call's response, the line with the call's own id — so
/// each call gets its own response, sealed, and the next call is not handed
/// the previous one's. Under `--upstream-self-contained` the awk server is
/// one process; whatever its figure is sealed as, it is the response with
/// the call's id that ends its span.
#[test]
fn the_binary_relays_what_a_server_sends_in_the_middle_of_a_call() {
    let mut g = Gateway::spawn(
        "chatty",
        &awk_server(TALKS_MID_CALL).iter().map(String::as_str).collect::<Vec<_>>(),
        &["--upstream-self-contained"],
    );
    g.send(&json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{}}}));
    assert_eq!(g.next("the initialize response")["id"], json!(1));
    g.send(&json!({"jsonrpc":"2.0","method":"notifications/initialized"}));
    for id in 2..5 {
        g.send(&call(id));
        let progress = g.next(&format!("call {id}'s progress notification"));
        assert_eq!(progress["method"], json!("notifications/progress"), "{progress}");
        assert_eq!(progress["params"]["progressToken"], json!(id), "{progress}");
        let ask = g.next(&format!("call {id}'s sampling request"));
        assert_eq!(ask["method"], json!("sampling/createMessage"), "{ask}");
        let ask_id = ask["id"].clone();
        assert_eq!(ask_id, json!(format!("s{id}")), "{ask}");
        g.send(&json!({"jsonrpc":"2.0","id":ask_id,"result":{"role":"assistant","content":{"type":"text","text":"ok"}}}));
        let r = g.next(&format!("the response to call {id}"));
        assert_eq!(r["id"], json!(id), "each call gets its own response: {r}");
        assert_eq!(r["result"]["content"][0]["text"], json!("sampled"), "the server got the client's answer mid-call: {r}");
        let meta = &r["result"]["_meta"]["science.transaction/jcp"];
        assert!(meta["receipt_b64"].is_string(), "{r}");
        let because = meta["energy_partial_because"].as_array().cloned().unwrap_or_default();
        assert!(!because.contains(&json!("reply-not-the-response")), "the span ended at the call's own response: {meta}");
        eprintln!("stdio binary: chatty call {id}: cost_uj {} [{}] domain {} because {because:?}", meta["cost_uj"], meta["energy_provenance"], meta["energy_domain"]);
    }
}

/// An MCP-ish server in awk that, inside each `tools/call`, tells the client
/// it has started (a `notifications/progress`) and then never answers: the
/// call's work is under way, and the gateway is waiting on it. Every other
/// request gets a plain response; notifications get nothing.
const STARTS_AND_HOLDS: &str = r#"/"method":"tools\/call"/ {
  id = "null"
  if (match($0, /"id":[0-9]+/)) id = substr($0, RSTART + 5, RLENGTH - 5)
  printf "{\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progressToken\":%s,\"progress\":1}}\n", id
  fflush()
  next
}
/"id"/ {
  id = "null"
  if (match($0, /"id":[0-9]+/)) id = substr($0, RSTART + 5, RLENGTH - 5)
  printf "{\"jsonrpc\":\"2.0\",\"id\":%s,\"result\":{\"protocolVersion\":\"2025-06-18\",\"capabilities\":{}}}\n", id
  fflush()
}"#;

/// Fault P2, in the binary: a gateway killed while a call's work was under
/// way lost the call's debit. The binary wrote a call's receipt to its
/// stream only after the relay returned, and the relay recorded nothing
/// before the work ran, so the stream a restart replayed held nothing of the
/// interrupted call and the grant came back whole (SPEC §4.11: the decision
/// is recorded before the action runs). The call's hold is now written and
/// synced before the call goes to the server (§8.1), so the restart holds
/// what the grant had left, refuses the next call, and names the
/// interrupted one as unresolved.
#[test]
fn p2_the_binary_killed_mid_call_loses_its_debit() {
    const BUDGET: u64 = 1_000_000_000;
    let holds = awk_server(STARTS_AND_HOLDS);
    let mut g = Gateway::spawn("killed", &holds.iter().map(String::as_str).collect::<Vec<_>>(), &[]);
    g.send(&json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{}}}));
    assert_eq!(g.next("the initialize response")["id"], json!(1));
    g.send(&json!({"jsonrpc":"2.0","method":"notifications/initialized"}));
    g.send(&call(2));
    let started = g.next("call 2's progress notification");
    assert_eq!(started["method"], json!("notifications/progress"), "{started}");
    // The server has the call and its work is under way: the gateway dies.
    g.child.kill().expect("kill the gateway");
    let _ = g.child.wait();
    drop(g);

    let answers = awk_server(ANSWERS_REQUESTS_ONLY);
    let mut g = Gateway::restart("killed", &answers.iter().map(String::as_str).collect::<Vec<_>>(), &[]);
    g.send(&json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{}}}));
    assert_eq!(g.next("the initialize response after the restart")["id"], json!(1));
    g.send(&json!({"jsonrpc":"2.0","method":"notifications/initialized"}));
    g.send(&call(3));
    let r = g.next("the response to call 3");
    assert_eq!(r["id"], json!(3), "{r}");
    // A refusal means the killed call's debit left too little for this one.
    let refused = r.get("error").is_some();
    let meta = &r["result"]["_meta"]["science.transaction/jcp"];
    let cost = meta["cost_uj"].as_u64().unwrap_or(0);
    let remaining = meta["budget_remaining_uj"].as_u64().unwrap_or(0);
    assert!(refused || meta.is_object(), "call 3 is answered or refused: {r}");
    assert!(
        refused || remaining + cost < BUDGET,
        "fault P2: the killed call's debit is gone: after the restart the grant held {remaining} + {cost} µJ of \
         {BUDGET}, all of it"
    );
    // The held call took all the grant had left, so the next is refused, and
    // the restart named the call it found in flight.
    assert!(r["error"]["message"].as_str().is_some_and(|m| m.contains("budget-exceeded")), "{r}");
    g.child.kill().expect("stop the gateway");
    let _ = g.child.wait();
    let stderr = g.stderr.recv_timeout(Duration::from_secs(5)).unwrap_or_default();
    assert!(stderr.contains("jcp-gateway: unresolved: the call held by receipt"), "{stderr}");
}

/// Whether `child` exits by `deadline`, and how: `None` if it is still running.
fn exits_within(child: &mut Child, deadline: Duration) -> Option<std::process::ExitStatus> {
    let start = std::time::Instant::now();
    loop {
        if let Ok(Some(status)) = child.try_wait() {
            return Some(status);
        }
        if start.elapsed() > deadline {
            return None;
        }
        std::thread::sleep(Duration::from_millis(50));
    }
}

/// SPEC §4.11, one writer per record: a second gateway started on a receipt
/// stream a running gateway holds refuses to serve, exit 2. Serving, the two
/// would each replay the same stream, each hold what the grant had left, and
/// each append, spending the budget twice on a forked chain. The first keeps
/// serving, and the claim goes with it: a gateway started on the file after
/// the first is killed serves.
#[test]
fn a_second_gateway_on_one_receipt_stream_refuses_to_serve() {
    let initialize = json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{}}});
    let answers = awk_server(ANSWERS_REQUESTS_ONLY);
    let server: Vec<&str> = answers.iter().map(String::as_str).collect();
    let mut first = Gateway::spawn("one-writer", &server, &[]);
    first.send(&initialize);
    // Answering, the first gateway holds its record: it takes it before it
    // reads a line.
    assert_eq!(first.next("the first gateway's initialize response")["id"], json!(1));

    let mut second = Gateway::restart("one-writer", &server, &[]);
    let Some(status) = exits_within(&mut second.child, LINE_DEADLINE) else {
        let _ = second.child.kill();
        let _ = second.child.wait();
        let stderr = second.stderr.recv_timeout(Duration::from_secs(5)).unwrap_or_default();
        panic!("a second gateway is serving on a receipt stream the first holds; its stderr:\n{stderr}");
    };
    let stderr = second.stderr.recv_timeout(Duration::from_secs(5)).unwrap_or_default();
    assert_eq!(status.code(), Some(2), "{stderr}");
    assert!(stderr.contains("refusing to serve: another process holds"), "{stderr}");
    drop(second);

    first.send(&json!({"jsonrpc":"2.0","method":"notifications/initialized"}));
    first.send(&call(2));
    let r = first.next("the first gateway's response to call 2");
    assert_eq!(r["id"], json!(2), "{r}");
    assert!(r["result"]["_meta"]["science.transaction/jcp"]["receipt_b64"].is_string(), "{r}");
    first.child.kill().expect("kill the first gateway");
    let _ = first.child.wait();
    drop(first);

    let mut after = Gateway::restart("one-writer", &server, &[]);
    after.send(&initialize);
    assert_eq!(after.next("the initialize response once the first gateway is gone")["id"], json!(1));
}