subako-function-sdk 0.2.2

use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use subako_function_sdk::{Content, Input, Output, Server, ServerInfo, ToolError, execute};

#[derive(Deserialize)]
struct Greet {
    name: String,
}

#[derive(Serialize)]
struct Greeting {
    text: String,
}

fn greet(tool: &str, input: Input) -> Result<Output, ToolError> {
    match tool {
        "greet" => {
            let greet: Greet = input.parse()?;
            Output::json(&Greeting {
                text: format!("hello {}", greet.name),
            })
        }
        other => Err(ToolError::unknown_tool(other)),
    }
}

/// A server listing one tool, whose calls the function it holds answers.
struct Serving<F>(F);

impl<F: Fn(&str, Input) -> Result<Output, ToolError>> Server for Serving<F> {
    fn info(&self) -> ServerInfo<'_> {
        ServerInfo {
            name: "hello",
            version: "1.2.3",
        }
    }

    fn tools(&self) -> Vec<Value> {
        vec![json!({"name": "greet", "inputSchema": {"type": "object"}})]
    }

    fn call(&self, tool: &str, input: Input) -> Result<Output, ToolError> {
        (self.0)(tool, input)
    }
}

fn meta() -> Value {
    json!({
        "io.modelcontextprotocol/protocolVersion": "2026-07-28",
        "io.modelcontextprotocol/clientCapabilities": {},
    })
}

/// A request as the worker writes it, `_meta` and all.
fn request(id: Value, method: &str, mut params: Value) -> Value {
    params["_meta"] = meta();
    json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params})
}

fn call(tool: &str, arguments: Value) -> Value {
    request(
        json!(1),
        "tools/call",
        json!({"name": tool, "arguments": arguments}),
    )
}

/// Every message the guest wrote with `stdin` on stdin, each line ending in
/// its newline.
fn written(stdin: &[u8], route: impl Fn(&str, Input) -> Result<Output, ToolError>) -> Vec<Value> {
    let (stdout, code) = execute(stdin, &Serving(route));
    assert_eq!(code, 0);
    let stdout = String::from_utf8(stdout).unwrap();
    assert!(stdout.is_empty() || stdout.ends_with('\n'), "{stdout}");
    stdout
        .lines()
        .map(|line| {
            let message: Value = serde_json::from_str(line).unwrap();
            assert_eq!(message["jsonrpc"], "2.0");
            message
        })
        .collect()
}

fn serve_lines(
    messages: &[Value],
    route: impl Fn(&str, Input) -> Result<Output, ToolError>,
) -> Vec<Value> {
    let stdin: String = messages
        .iter()
        .map(|message| format!("{message}\n"))
        .collect();
    written(stdin.as_bytes(), route)
}

/// The one message the guest answered `message` with.
fn response(message: Value, route: impl Fn(&str, Input) -> Result<Output, ToolError>) -> Value {
    let mut written = serve_lines(&[message], route);
    assert_eq!(written.len(), 1, "{written:?}");
    written.remove(0)
}

/// The result the guest answered `message` with, as id 1, less what every
/// result carries.
fn result(message: Value, route: impl Fn(&str, Input) -> Result<Output, ToolError>) -> Value {
    let response = response(message, route);
    assert_eq!(response["id"], 1);
    let mut result = response["result"].clone();
    let result = result.as_object_mut().unwrap();
    assert_eq!(result.remove("resultType"), Some(json!("complete")));
    assert_eq!(
        result.remove("_meta"),
        Some(json!({"io.modelcontextprotocol/serverInfo": {"name": "hello", "version": "1.2.3"}}))
    );
    Value::Object(result.clone())
}

/// The JSON-RPC error the guest answered `message` with, never having
/// routed.
fn refusal(message: Value) -> Value {
    let response = response(message, |_, _| panic!("routed a refused request"));
    assert!(response.get("result").is_none(), "{response}");
    response
}

/// The JSON-RPC error the guest answered the one line on `stdin` with.
fn line_refusal(stdin: &[u8]) -> Value {
    let mut written = written(stdin, |_, _| panic!("routed a refused request"));
    assert_eq!(written.len(), 1, "{written:?}");
    written.remove(0)
}

fn failure_text(result: &Value) -> String {
    assert_eq!(result["isError"], true);
    result["content"][0]["text"].as_str().unwrap().to_string()
}

#[test]
fn discover_names_the_revision_and_the_tools_capability() {
    assert_eq!(
        result(request(json!(1), "server/discover", json!({})), greet),
        json!({"supportedVersions": ["2026-07-28"], "capabilities": {"tools": {}}})
    );
}

#[test]
fn tools_list_holds_every_tool_on_one_page() {
    assert_eq!(
        result(request(json!(1), "tools/list", json!({})), greet),
        json!({
            "tools": [{"name": "greet", "inputSchema": {"type": "object"}}],
            "ttlMs": 86_400_000,
            "cacheScope": "public",
        })
    );
    let next = refusal(request(json!(1), "tools/list", json!({"cursor": ""})));
    assert_eq!(next["error"]["code"], -32602);
}

#[test]
fn each_request_is_answered_in_turn() {
    let written = serve_lines(
        &[
            request(json!(0), "server/discover", json!({})),
            json!({"jsonrpc": "2.0", "method": "notifications/cancelled", "params": {}}),
            call("greet", json!({"name": "ada"})),
            call("greet", json!({"name": "grace"})),
        ],
        greet,
    );
    let ids: Vec<&Value> = written.iter().map(|message| &message["id"]).collect();
    assert_eq!(ids, [&json!(0), &json!(1), &json!(1)]);
    assert_eq!(
        written[1]["result"]["structuredContent"],
        json!({"text": "hello ada"})
    );
    assert_eq!(
        written[2]["result"]["structuredContent"],
        json!({"text": "hello grace"})
    );
}

#[test]
fn a_known_tool_answers_the_request_with_its_result() {
    assert_eq!(
        result(call("greet", json!({"name": "ada"})), greet),
        json!({
            "content": [{"type": "text", "text": r#"{"text":"hello ada"}"#}],
            "structuredContent": {"text": "hello ada"},
        })
    );
}

#[test]
fn the_response_echoes_the_request_id_as_it_came() {
    let request = request(
        json!("call-7"),
        "tools/call",
        json!({"name": "greet", "arguments": {"name": "ada"}}),
    );
    assert_eq!(response(request, greet)["id"], "call-7");
}

#[test]
fn content_alone_carries_no_json_beside_it() {
    let out = result(call("any", json!({})), |_, _| {
        Ok(Output::content(vec![
            Content::text("look"),
            Content::Image {
                mime_type: "image/png".to_string(),
                data: "iVBOR".to_string(),
            },
        ]))
    });
    assert_eq!(
        out,
        json!({"content": [
            {"type": "text", "text": "look"},
            {"type": "image", "mimeType": "image/png", "data": "iVBOR"},
        ]})
    );
}

#[test]
fn an_image_file_is_a_link_to_its_place_in_the_volume() {
    let out = result(call("any", json!({})), |_, _| {
        Ok(Output::content(vec![
            Content::image_file("/volume/data/my plots/a+b.png"),
            Content::image_file("/tmp/c.png"),
        ]))
    });
    assert_eq!(
        out,
        json!({"content": [
            {"type": "resource_link", "uri": "volume:///data/my%20plots/a%2Bb.png", "name": "a+b.png"},
            {"type": "resource_link", "uri": "file:///tmp/c.png", "name": "c.png"},
        ]})
    );
}

#[test]
fn an_unknown_tool_is_a_tool_error() {
    let out = result(call("nope", json!({})), greet);
    assert_eq!(failure_text(&out), "unknown tool: nope");
}

#[test]
fn arguments_of_the_wrong_shape_name_the_problem() {
    let out = result(call("greet", json!({"name": 3})), greet);
    assert!(failure_text(&out).starts_with("invalid input"));
}

#[test]
fn absent_arguments_are_an_empty_object() {
    let request = request(json!(1), "tools/call", json!({"name": "echo"}));
    let out = result(request, |_, input| Output::json(&input.parse::<Value>()?));
    assert_eq!(out["structuredContent"], json!({}));
}

#[test]
fn members_the_guest_does_not_know_are_read_past() {
    let mut request = call("greet", json!({"name": "ada"}));
    request["params"]["_meta"]["progressToken"] = json!("p");
    request["params"]["_meta"]["io.modelcontextprotocol/clientInfo"] =
        json!({"name": "subako", "version": "1"});
    request["params"]["from_the_future"] = json!(true);
    request["from_the_future"] = json!(true);
    let out = result(request, greet);
    assert_eq!(out["structuredContent"], json!({"text": "hello ada"}));
}

#[test]
fn a_request_without_its_meta_has_invalid_params() {
    let bare = |method: &str| json!({"jsonrpc": "2.0", "id": 4, "method": method, "params": {}});
    for request in [
        bare("tools/call"),
        bare("initialize"),
        json!({"jsonrpc": "2.0", "id": 4, "method": "tools/list"}),
        json!({
            "jsonrpc": "2.0",
            "id": 4,
            "method": "tools/list",
            "params": {"_meta": {"io.modelcontextprotocol/protocolVersion": "2026-07-28"}},
        }),
    ] {
        let out = refusal(request.clone());
        assert_eq!(out["id"], 4, "{request}");
        assert_eq!(out["error"]["code"], -32602, "{request}");
        assert!(
            out["error"]["message"]
                .as_str()
                .unwrap()
                .contains("2026-07-28"),
            "{out}"
        );
    }
}

#[test]
fn another_revision_is_unsupported() {
    let mut request = call("greet", json!({}));
    request["params"]["_meta"]["io.modelcontextprotocol/protocolVersion"] = json!("2099-01-01");
    let out = refusal(request);
    assert_eq!(
        out["error"],
        json!({
            "code": -32022,
            "message": "Unsupported protocol version",
            "data": {"supported": ["2026-07-28"], "requested": "2099-01-01"},
        })
    );
}

#[test]
fn notifications_and_responses_go_unanswered() {
    assert!(
        serve_lines(
            &[
                json!({"jsonrpc": "2.0", "method": "notifications/cancelled", "params": {}}),
                json!({"jsonrpc": "2.0", "id": 3, "result": {}}),
                json!({"jsonrpc": "2.0", "id": 3, "error": {"code": 1, "message": "no"}}),
            ],
            |_, _| panic!("routed a notification")
        )
        .is_empty()
    );
}

#[test]
fn an_empty_stdin_is_answered_with_nothing() {
    assert!(written(b"", greet).is_empty());
}

#[test]
fn a_line_that_is_not_json_is_a_parse_error_and_the_next_is_still_served() {
    let stdin = format!("not json\n{}\n", call("greet", json!({"name": "ada"})));
    let written = written(stdin.as_bytes(), greet);
    assert_eq!(written[0]["id"], Value::Null);
    assert_eq!(written[0]["error"]["code"], -32700);
    assert_eq!(
        written[1]["result"]["structuredContent"],
        json!({"text": "hello ada"})
    );
    assert_eq!(line_refusal(b"\n")["error"]["code"], -32700, "a blank line");
}

#[test]
fn json_that_is_no_message_is_an_invalid_request() {
    for (line, id) in [
        (json!([1, 2]), Value::Null),
        (json!({"id": 4, "method": "tools/list"}), Value::Null),
        (json!({"jsonrpc": "2.0"}), Value::Null),
        (json!({"jsonrpc": "2.0", "id": 4}), json!(4)),
    ] {
        let out = line_refusal(format!("{line}\n").as_bytes());
        assert_eq!(out["id"], id, "{line}");
        assert_eq!(out["error"]["code"], -32600, "{line}");
    }
}

#[test]
fn another_method_is_not_found() {
    let out = refusal(request(json!(4), "resources/list", json!({})));
    assert_eq!(out["id"], 4);
    assert_eq!(out["error"]["code"], -32601);
}

#[test]
fn another_json_rpc_version_is_an_invalid_request() {
    let mut request = call("greet", json!({}));
    request["jsonrpc"] = json!("1.0");
    request["id"] = json!(4);
    let out = refusal(request);
    assert_eq!(out["id"], 4);
    assert_eq!(out["error"]["code"], -32600);
}

#[test]
fn a_call_naming_no_tool_has_invalid_params() {
    let out = refusal(request(json!(4), "tools/call", json!({"arguments": {}})));
    assert_eq!(out["id"], 4);
    assert_eq!(out["error"]["code"], -32602);
}

#[test]
fn output_that_json_cannot_hold_is_a_tool_error() {
    // JSON object keys are strings, so a map keyed by pairs has no JSON form.
    let unjsonable = std::collections::BTreeMap::from([((1, 2), 0)]);
    let error = Output::json(&unjsonable).err().unwrap();
    assert!(
        error.message().starts_with("output is not JSON"),
        "{error:?}"
    );
}

#[test]
fn a_tool_error_takes_a_message_of_either_string_type() {
    assert_eq!(ToolError::from("no").message(), "no");
    assert_eq!(ToolError::from("no".to_string()).message(), "no");
}