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()
}"#;

/// 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 receipts = std::path::Path::new(env!("CARGO_TARGET_TMPDIR"))
            .join(format!("jcp-gateway-{name}-{}.cose", std::process::id()));
        let _ = std::fs::remove_file(&receipts);
        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"]);
    }
}