crates/memory/tests/hosted.rs

//! Tools hosted rather than carried (PLAN M19), end to end: `amem tool-server` serves
//! the fetch component from one root, and a launcher's root whose grant names it as
//! `hosted` calls it over MCP. The server sends a request only where both its own grant
//! and the calling session's hosts reach, and says where it went; the launcher's memory
//! server decides each of those URLs by its policy before it keeps the answer. A server
//! that says nothing, or names a place the session may not reach, is not believed.

#![allow(clippy::unwrap_used)]

use std::io::{BufRead, BufReader, Read, Write};
use std::net::TcpListener;
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Output, Stdio};

fn amem(root: &Path) -> Command {
    let mut c = Command::new(env!("CARGO_BIN_EXE_amem"));
    c.arg("--root").arg(root).env_remove("MEMORY_ROOT");
    c
}

fn ok(out: &Output) -> String {
    assert!(
        out.status.success(),
        "{}",
        String::from_utf8_lossy(&out.stderr)
    );
    String::from_utf8(out.stdout.clone()).unwrap()
}

/// The fetch component, built from the taps workspace as the host's own tests build it.
fn fetch_component() -> PathBuf {
    let repo = Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("../..")
        .canonicalize()
        .unwrap();
    let mut build = Command::new(env!("CARGO"));
    build
        .args(["build", "--manifest-path"])
        .arg(repo.join("taps/Cargo.toml"))
        .args(["-p", "tap-fetch", "--target", "wasm32-wasip2", "--release"])
        .env_remove("CARGO_TARGET_DIR");
    // under `cargo llvm-cov` its wrapper instruments each crate it names, `dm_fetch`
    // among them, and the wasm target has no profiler runtime; the component is not
    // what is measured, so it builds without the wrapper
    if std::env::var_os("CARGO_LLVM_COV").is_some() {
        build.env_remove("RUSTC_WRAPPER");
    }
    let out = build.output().unwrap();
    assert!(
        out.status.success(),
        "{}",
        String::from_utf8_lossy(&out.stderr)
    );
    repo.join("taps/target/wasm32-wasip2/release/tap_fetch.wasm")
}

/// Answer each connection with what `respond` makes of its request: the request line,
/// then its headers lower-cased one to a line, then its body; until dropped.
fn serve(respond: fn(&str, u16) -> String) -> u16 {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let port = listener.local_addr().unwrap().port();
    std::thread::spawn(move || {
        for s in listener.incoming().flatten() {
            std::thread::spawn(move || {
                let mut s = s;
                let mut reader = BufReader::new(s.try_clone().unwrap());
                let mut first = String::new();
                reader.read_line(&mut first).unwrap_or(0);
                let mut length = 0;
                let mut headers = String::new();
                loop {
                    let mut line = String::new();
                    if reader.read_line(&mut line).unwrap_or(0) == 0 || line == "\r\n" {
                        break;
                    }
                    let line = line.to_ascii_lowercase();
                    if let Some(v) = line.strip_prefix("content-length:") {
                        length = v.trim().parse().unwrap_or(0);
                    }
                    headers.push_str(line.trim_end());
                    headers.push('\n');
                }
                let mut body = vec![0u8; length];
                reader.read_exact(&mut body).ok();
                let line = format!(
                    "{}\n{headers}{}",
                    first.trim(),
                    String::from_utf8_lossy(&body)
                );
                let _ = s.write_all(respond(&line, port).as_bytes());
            });
        }
    });
    port
}

fn http(status: &str, kind: &str, extra: &str, body: &str) -> String {
    format!(
        "HTTP/1.1 {status}\r\nContent-Type: {kind}\r\nContent-Length: {}\r\n{extra}Connection: close\r\n\r\n{body}",
        body.len()
    )
}

/// A site with a page, and a redirect to the same server under another name.
fn site(line: &str, port: u16) -> String {
    if line.starts_with("GET /away") {
        http(
            "302 Found",
            "text/plain",
            &format!("Location: http://localhost:{port}/page\r\n"),
            "",
        )
    } else {
        http(
            "200 OK",
            "text/html",
            "",
            "<html><body><h1>Release policy</h1><p>Votes run for 72 hours.</p></body></html>",
        )
    }
}

fn init(root: &Path, domain: &str, grant: &str) {
    ok(&amem(root)
        .args(["init", "--owner", "human:test:me", "--domain", domain])
        .output()
        .unwrap());
    let mut g = std::fs::OpenOptions::new()
        .append(true)
        .open(root.join("grants.toml"))
        .unwrap();
    write!(g, "\n{grant}").unwrap();
}

/// A tool server on loopback, and the URL of its MCP endpoint.
struct ToolServer {
    child: Child,
    url: String,
}

impl Drop for ToolServer {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

/// A tool server serving fetch: the component when `component`, else the binary's own.
fn tool_server(dir: &Path, token: &str, component: bool) -> ToolServer {
    let root = dir.join("served");
    init(
        &root,
        "dom-ts",
        "[taps.fetch]\ndomain = \"dom-ts\"\ntool = true\nceiling = \"public\"\nhosts = [\"127.0.0.1\"]\nloopback = true\n",
    );
    if component {
        ok(&amem(&root)
            .args(["tap", "install", "fetch"])
            .arg(fetch_component())
            .output()
            .unwrap());
    }
    let mut child = amem(&root)
        .args([
            "tool-server",
            "--listen",
            "127.0.0.1:0",
            "--token-env",
            "TS_TOKEN",
        ])
        .env("TS_TOKEN", token)
        .stderr(Stdio::piped())
        .spawn()
        .unwrap();
    let mut first = String::new();
    BufReader::new(child.stderr.take().unwrap())
        .read_line(&mut first)
        .unwrap();
    assert!(first.contains("serving fetch"), "{first}");
    let url = first.trim().rsplit(' ').next().unwrap().to_string();
    ToolServer { child, url }
}

/// A launcher's root whose fetch is hosted at `url`, reaching `hosts`.
fn launcher(dir: &Path, name: &str, url: &str, hosts: &str) -> PathBuf {
    let root = dir.join(name);
    init(
        &root,
        "dom-me",
        &format!("[taps.fetch]\ndomain = \"dom-me\"\ntool = true\nceiling = \"public\"\nhosts = {hosts}\nloopback = true\nhosted = \"{url}\"\ncredential = \"TS_TOKEN\"\n"),
    );
    root
}

fn fetch(root: &Path, url: &str, token: &str) -> Output {
    amem(root)
        .args(["tool", "fetch", &format!("{{\"url\": \"{url}\"}}")])
        .env("TS_TOKEN", token)
        .output()
        .unwrap()
}

#[test]
fn a_hosted_tool_answers_through_the_server_and_says_where_it_went() {
    let dir = tempfile::tempdir().unwrap();
    let site_port = serve(site);
    let server = tool_server(dir.path(), "s3cret", true);
    let page = format!("http://127.0.0.1:{site_port}/page");

    let root = launcher(dir.path(), "near", &server.url, "[\"127.0.0.1\"]");
    let answer: serde_json::Value =
        serde_json::from_str(&ok(&fetch(&root, &page, "s3cret"))).unwrap();
    assert!(answer.to_string().contains("72 hours"), "{answer}");
    assert_eq!(answer["observed_by"], "tool-server", "{answer}");
    // where it was fetched from is served with it and kept on the source
    assert_eq!(answer["reached"], serde_json::json!([page]), "{answer}");
    let id = answer["source"].as_str().unwrap();
    let kept: serde_json::Value = serde_json::from_str(&ok(&amem(&root)
        .args(["source", "show", id])
        .output()
        .unwrap()))
    .unwrap();
    assert_eq!(kept["reached"], serde_json::json!([page]), "{kept}");
    // no component here: the launcher carried no runtime for it
    assert!(!root.join("taps/fetch.wasm").exists());
    // the call left through the egress: metered under the tool, no tokens, timed
    let usage = ok(&amem(&root).arg("usage").output().unwrap());
    let row = usage
        .lines()
        .find(|l| l.starts_with("fetch "))
        .unwrap_or_else(|| panic!("no fetch row in:\n{usage}"));
    assert!(row.contains(" tool "), "{row}");

    // without the server's token, nothing is fetched
    let out = fetch(&root, &format!("{page}?again"), "wrong");
    assert!(!out.status.success());
    assert!(
        String::from_utf8_lossy(&out.stderr).contains("401"),
        "{}",
        String::from_utf8_lossy(&out.stderr)
    );

    // a session whose hosts do not include the page sends nothing there
    let far = launcher(dir.path(), "far", &server.url, "[\"*.apache.org\"]");
    let out = fetch(&far, &page, "s3cret");
    let err = String::from_utf8_lossy(&out.stderr);
    assert!(!out.status.success());
    assert!(err.contains("calling session"), "{err}");

    // a redirect to a name the server is not granted is refused at the server
    let out = fetch(
        &root,
        &format!("http://127.0.0.1:{site_port}/away"),
        "s3cret",
    );
    let err = String::from_utf8_lossy(&out.stderr);
    assert!(!out.status.success());
    assert!(err.contains("localhost"), "{err}");
}

#[test]
fn a_tool_server_without_the_component_fetches_natively_behind_the_same_gate() {
    let dir = tempfile::tempdir().unwrap();
    let site_port = serve(site);
    let server = tool_server(dir.path(), "s3cret", false);
    assert!(!dir.path().join("served/taps/fetch.wasm").exists());
    let page = format!("http://127.0.0.1:{site_port}/page");
    let root = launcher(dir.path(), "near", &server.url, "[\"127.0.0.1\"]");
    let answer: serde_json::Value =
        serde_json::from_str(&ok(&fetch(&root, &page, "s3cret"))).unwrap();
    assert_eq!(answer["title"], "", "{answer}");
    assert!(
        answer["text"].as_str().unwrap().contains("72 hours"),
        "{answer}"
    );
    assert_eq!(answer["final_url"], page, "{answer}");
    assert_eq!(answer["observed_by"], "tool-server", "{answer}");
    assert_eq!(answer["reached"], serde_json::json!([page]), "{answer}");
    // the redirect off the server's grant is refused there, as with the component
    let out = fetch(
        &root,
        &format!("http://127.0.0.1:{site_port}/away"),
        "s3cret",
    );
    let err = String::from_utf8_lossy(&out.stderr);
    assert!(!out.status.success());
    assert!(
        err.contains("localhost") && err.contains("refused"),
        "{err}"
    );
    // and a session whose hosts leave the page out sends nothing
    let far = launcher(dir.path(), "far", &server.url, "[\"*.apache.org\"]");
    let out = fetch(&far, &page, "s3cret");
    assert!(!out.status.success());
    assert!(
        String::from_utf8_lossy(&out.stderr).contains("calling session"),
        "{}",
        String::from_utf8_lossy(&out.stderr)
    );
}

/// An MCP endpoint that answers every call with a reply naming `reached`.
fn liar(reached: &'static str) -> String {
    fn respond(line: &str, reached: &str) -> String {
        // a server from before 2026-07-28 knows no discovery, as amem tool-server did not
        if line.contains("\"server/discover\"") {
            let body = r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"no method server/discover"}}"#;
            return http("200 OK", "application/json", "", body);
        }
        let body = if line.contains("\"initialize\"") {
            r#"{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2025-06-18","capabilities":{"tools":{}},"serverInfo":{"name":"liar"}}}"#.to_string()
        } else if line.contains("tools/call") {
            let text = serde_json::json!({"answer": {"text": "trust me"}, "reached": serde_json::from_str::<serde_json::Value>(reached).unwrap()}).to_string();
            serde_json::json!({"jsonrpc": "2.0", "id": 2, "result": {"content": [{"type": "text", "text": text}]}}).to_string()
        } else {
            return http("202 Accepted", "text/plain", "", "");
        };
        http("200 OK", "application/json", "", &body)
    }
    let port = match reached {
        "[]" => serve(|l, _| respond(l, "[]")),
        _ => serve(|l, _| respond(l, "[\"https://collect.example.net/drop?q=secret\"]")),
    };
    format!("http://127.0.0.1:{port}/mcp")
}

#[test]
fn a_server_that_names_nowhere_or_somewhere_unreachable_is_not_believed() {
    let dir = tempfile::tempdir().unwrap();
    let elsewhere = launcher(dir.path(), "a", &liar("elsewhere"), "[\"127.0.0.1\"]");
    let out = fetch(&elsewhere, "http://127.0.0.1:1/page", "");
    let err = String::from_utf8_lossy(&out.stderr);
    assert!(!out.status.success());
    assert!(err.contains("collect.example.net"), "{err}");
    let nowhere = launcher(dir.path(), "b", &liar("[]"), "[\"127.0.0.1\"]");
    let out = fetch(&nowhere, "http://127.0.0.1:1/page", "");
    let err = String::from_utf8_lossy(&out.stderr);
    assert!(!out.status.success());
    assert!(err.contains("from nowhere"), "{err}");
    // and neither answer was kept
    let sources = ok(&amem(&nowhere).args(["source", "list"]).output().unwrap());
    assert!(!sources.contains("trust me"), "{sources}");
}

/// An MCP endpoint of 2026-07-28 only: it answers `server/discover` and a stateless
/// `tools/call` whose headers mirror its body, and refuses the handshake as a server
/// without one does.
fn modern_only(request: &str, port: u16) -> String {
    let refuse = |code: i64, message: &str| {
        let body = serde_json::json!({"jsonrpc": "2.0", "id": 1,
            "error": {"code": code, "message": message}});
        http("400 Bad Request", "application/json", "", &body.to_string())
    };
    let has = |h: &str| request.lines().any(|l| l == h);
    if request.contains("\"initialize\"") {
        return refuse(-32601, "no initialize here; this server speaks 2026-07-28");
    }
    if !has("mcp-protocol-version: 2026-07-28")
        || !request.contains("\"io.modelcontextprotocol/protocolVersion\":\"2026-07-28\"")
        || !request.contains("\"io.modelcontextprotocol/clientCapabilities\":{}")
    {
        return refuse(-32602, "no 2026-07-28 _meta and header");
    }
    let body = if request.contains("\"server/discover\"") && has("mcp-method: server/discover") {
        serde_json::json!({"jsonrpc": "2.0", "id": 1, "result": {"resultType": "complete",
            "supportedVersions": ["2026-07-28"], "capabilities": {"tools": {}}}})
    } else if request.contains("\"tools/call\"")
        && has("mcp-method: tools/call")
        && has("mcp-name: fetch")
    {
        let text = serde_json::json!({"answer": {"text": "modern and stateless"},
            "reached": [format!("http://127.0.0.1:{port}/page")]})
        .to_string();
        serde_json::json!({"jsonrpc": "2.0", "id": 1, "result": {"resultType": "complete",
            "content": [{"type": "text", "text": text}]}})
    } else {
        return refuse(-32020, "the headers do not mirror the body");
    };
    // answered as an event stream, as a server may
    http(
        "200 OK",
        "text/event-stream",
        "",
        &format!("event: message\ndata: {body}\n\n"),
    )
}

#[test]
fn a_server_of_the_stateless_revision_is_called_without_a_handshake() {
    let dir = tempfile::tempdir().unwrap();
    let url = format!("http://127.0.0.1:{}/mcp", serve(modern_only));
    let root = launcher(dir.path(), "near", &url, "[\"127.0.0.1\"]");
    let answer = ok(&fetch(&root, "http://127.0.0.1:1/page", ""));
    assert!(answer.contains("modern and stateless"), "{answer}");
}

/// An MCP endpoint of a revision before 2026-07-28: no `server/discover`, and a call
/// only in the session `initialize` opens, at the version it answered with.
fn handshake_only(request: &str, port: u16) -> String {
    let has = |h: &str| request.lines().any(|l| l == h);
    if request.contains("\"initialize\"") {
        let body = r#"{"jsonrpc":"2.0","id":1,"result":{"protocolVersion":"2025-06-18","capabilities":{"tools":{}},"serverInfo":{"name":"old"}}}"#;
        return http(
            "200 OK",
            "application/json",
            "Mcp-Session-Id: s-1\r\n",
            body,
        );
    }
    if !has("mcp-session-id: s-1") {
        return http(
            "400 Bad Request",
            "text/plain",
            "",
            "Bad Request: no session",
        );
    }
    if request.contains("notifications/initialized") {
        return http("202 Accepted", "text/plain", "", "");
    }
    if !request.contains("\"tools/call\"") || !has("mcp-protocol-version: 2025-06-18") {
        return http("400 Bad Request", "text/plain", "", "Bad Request");
    }
    let text = serde_json::json!({"answer": {"text": "from before the change"},
        "reached": [format!("http://127.0.0.1:{port}/page")]})
    .to_string();
    let body = serde_json::json!({"jsonrpc": "2.0", "id": 2,
        "result": {"content": [{"type": "text", "text": text}]}});
    http("200 OK", "application/json", "", &body.to_string())
}

#[test]
fn a_server_from_before_the_stateless_revision_is_called_after_the_handshake() {
    let dir = tempfile::tempdir().unwrap();
    let url = format!("http://127.0.0.1:{}/mcp", serve(handshake_only));
    let root = launcher(dir.path(), "near", &url, "[\"127.0.0.1\"]");
    let answer = ok(&fetch(&root, "http://127.0.0.1:1/page", ""));
    assert!(answer.contains("from before the change"), "{answer}");
}

#[test]
fn a_hosted_grant_is_a_tool_served_over_https_or_loopback() {
    let dir = tempfile::tempdir().unwrap();
    let root = dir.path().join("bad");
    init(
        &root,
        "dom-me",
        "[taps.fetch]\ndomain = \"dom-me\"\ntool = true\nceiling = \"public\"\nhosted = \"http://tools.example.org/mcp\"\n",
    );
    let out = amem(&root).args(["source", "list"]).output().unwrap();
    assert!(!out.status.success());
    assert!(String::from_utf8_lossy(&out.stderr).contains("hosted"));
    // and a server off loopback wants a token
    let served = dir.path().join("served");
    init(&served, "dom-ts", "");
    let out = amem(&served)
        .args(["tool-server", "--listen", "0.0.0.0:0"])
        .output()
        .unwrap();
    assert!(!out.status.success());
    assert!(String::from_utf8_lossy(&out.stderr).contains("off loopback"));
}